
%% Comments 
% YBP project (MITgcm, Niagara)
%
% Kinetic energy and shear spectra from Niagara results with 
% WKB scaling and the depth stretched.
%
%
% Author: Ritabrata Thakur, 1 Oct 2021.
%
% Uses: waveno_only_spectra_shear_ybp.m for Pinkel estimate of shear
% and Gregg estimate is done here itself.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Mode of usage: fully automatic
%
% Last used: 28/3/22
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear

% update on 19 Jan 22: run_no should be 19,20,1,2,3,4 as the 153 and 264
% are not the runs that I had originally. 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% check run_logs_new.odp
run_no_array = {19,20,1,2,3,4};
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for jj=1:length(run_no_array)

run_no = run_no_array{jj};

%check for "run_no" variable in the excel sheet
if (run_no == 19 || run_no == 20)
 vert_levs = 109;
end
if (run_no == 1 || run_no == 2)
    vert_levs = 153;
end
if (run_no == 3 || run_no == 4)
    vert_levs = 264;
end

%if (run_no == 34 || run_no == 35)
%    vert_levs = 153;
%end

%if (run_no == 25 || run_no == 26)
%    vert_levs = 264;
%end

fprintf('Running vertical level %d\n',vert_levs);
for MMP_number=[1 3 4]
fprintf('Running MMP_number %d\n',MMP_number);
clearvars -except jj run_no vert_levs MMP_number run_no_array 
%% Reading variables
if (run_no == 19 || run_no == 20)
 [data,lat,lon,ix,jx,ncpath,XC,YC,RC,time_start,time_end,time_array,z_ip] = read_ncfiles(vert_levs,MMP_number,run_no);
end

if (run_no == 1 || run_no == 2 || run_no == 3 || run_no == 4)
 [data,lat,lon,ix,jx,ncpath,XC,YC,RC,time_start,time_end,time_array,z_ip] = read_ncfiles_km_runs(vert_levs,MMP_number,run_no);
end

if (MMP_number == 1)
 inertial_start = 25.05; % 90% of 27.84 hrs
 inertial_end   = 30.6;  % 110% of 27.84 hrs
end

if (MMP_number == 3)
 inertial_start = 22.34; % 90% of 24.83 hrs
 inertial_end   = 27.31; % 110% of 24.83 hrs
end

if (MMP_number == 4)
 inertial_start = 21.52; % 90% of 23.92 hrs
 inertial_end   = 26.31; % 110% of 23.92 hrs
end
 
%% N2 calculation for WKB scaling. Needs GSW toolbox

%Z = depth_array;
p = gsw_p_from_z(squeeze(RC),lat);

% Taking salt and salinity from the "all_ON" run
[SA, ~] = gsw_SA_from_SP(squeeze(data.salt(1,1,:,:)),p,lon,lat);  %converts practical salinity
%to absolute salinity which is used by the N2 routine. The flag of "1" if
%the measurement is in ocean. If it is used in land, set it to 0. do not
%put anything. it is automatically set
CT = gsw_CT_from_t(SA,squeeze(data.temp(1,1,:,:)),p); %converts in-situ temperature to
%conservative temperature as is required in GSW toolbox

clear N2 NN
for i=1:length(CT(1,:))
 [N2,p_mid] = gsw_Nsquared(SA(:,i),CT(:,i),p,lat);
 NN(i,:)=N2;  %this is the total Brunt-vaisala frequency squared stored
 % at a single location. It is probably best to continue with one location
 % as averaging over different locations would mean we would have to
 % calculate the wkb factor for each depth differently as the
 % stratification would vary according to location
end

% removing the values where there is unstable stratification
% because when we calculate Brunt Vaisala frequency, it would otherwise
% become complex
for i=1:length(NN(:,1))
    for j=1:length(NN(1,:))
      if(NN(i,j) <= 0)
      NN(i,j) = 0;
      end
    end
end

%sqrt of the N^2
BV = sqrt(NN);

clear BV_interp 
%interpolating to original depth levels. Remember that "p" and the
%depth_array are different 
for i=1:length(NN(:,1))
 BV_interp(i,:) = interp1(p_mid, BV(i,:), p, 'linear', 'extrap');
end

% choosing intermediate depth range. This can be used for both 109 level
% and 264 level. 
start_depth=find(abs(p_mid-80)<5,1); %taking approx 80m as the start depth to match MMP 
end_depth=find(abs(p_mid-1400)<50,1,'last');  %taking approx 1400m as the end depth to match MMP 

% Brunt Vaisala in the intermediate (MMP) depth range
depth_array_short = squeeze(RC(start_depth:end_depth));
BV_interp_short   = squeeze(BV_interp(time_start:time_end,start_depth:end_depth));
BV_interp_time_mean =  mean(BV_interp_short, 'omitnan');
BV_interp_time_mean = naninterp(BV_interp_time_mean);
NN_interp_time_mean = BV_interp_time_mean.^2;

N_0_mean = 0.0052; %used in literature = 3cph

wkb_factor_short = sqrt(BV_interp_time_mean./N_0_mean);

%% WKB scaling
% stretched depth from the chosen intermediate depth
clear z_stretched
z_stretched(1) = (-depth_array_short(1));
%for j=1:length(NN_interp)
 for ii=2:length(depth_array_short)
  z_stretched(ii) = z_stretched(1) + trapz(-depth_array_short(1:ii), squeeze(BV_interp_time_mean(1:ii)))./N_0_mean; 
 end
%end


% WKB like scaling of the horizontal velocities at the intermediate depth
% where the stratification does not vary abruptly
clear data_uscaled data_vscaled
%for i=1:length(data.U(:,1,1,1))
%    for j=1:length(data.U(1,:,1,1))
kk=1;
for k=time_start:time_end
  data_uscaled(:,kk) = squeeze(data.U(1,1,start_depth:end_depth,k))'./wkb_factor_short;
  data_vscaled(:,kk) = squeeze(data.V(1,1,start_depth:end_depth,k))'./wkb_factor_short;
  kk=kk+1;
end
%    end
%end

%% Stretched vertical coordinate 
% interpolating to stretched coordinates
% this array is needed to interpolate the velocities from unequally-spaced
% stretched coordinates to equally-spaced. See Leaman and Sanford pp 1976
z_stretched_equal_spaced = linspace(z_stretched(1), z_stretched(end), length(z_stretched(1:end)));
%z_unstretched_equal      = linspace(depth_array_short(1), depth_array_short(end), length(depth_array_short(1:end)));

% 26 Jan 2022
clear grid_old_x grid_new_x grid_old_y grid_new_y grid_old_z grid_new_z grid_old_time grid_new_time
[grid_old_z, grid_old_time] = ndgrid(squeeze(z_stretched), time_array);
[grid_new_z, grid_new_time] = ndgrid(squeeze(z_stretched_equal_spaced), time_array);

scaled_eq_z.U = interpn(grid_old_z, grid_old_time, data_uscaled, grid_new_z, grid_new_time);
scaled_eq_z.V = interpn(grid_old_z, grid_old_time, data_vscaled, grid_new_z, grid_new_time);

%% Filtering velocities in diurnal and near-inertial
% bands to calculate shear in those hourly bands. Near-inertial is the most
% energetic after eddies in a global-averaged sense, Jan 25, 2022. I also
% see that the semidiurnal band is strong here from n018 plot

for i=1:length(scaled_eq_z.U(:,1))
  scaled_eq_z.U_6_to_8hr(i,:)  = bandpass(squeeze(scaled_eq_z.U(i,:)),[1/8.20 1/5.80], 1); 
  scaled_eq_z.U_sem_diur(i,:)  = bandpass(squeeze(scaled_eq_z.U(i,:)),[1/13.5 1/11.5], 1); 
  scaled_eq_z.U_highpass(i,:)  = highpass(squeeze(scaled_eq_z.U(i,:)),1/11.5, 1);   
  scaled_eq_z.U_lowpass(i,:)   = lowpass(squeeze(scaled_eq_z.U(i,:)),1/11.5, 1);   
  scaled_eq_z.U_inertial(i,:)  = bandpass(squeeze(scaled_eq_z.U(i,:)),[1/inertial_end 1/inertial_start], 1); 
  scaled_eq_z.V_6_to_8hr(i,:)  = bandpass(squeeze(scaled_eq_z.V(i,:)),[1/8.20 1/5.80], 1); 
  scaled_eq_z.V_sem_diur(i,:)  = bandpass(squeeze(scaled_eq_z.V(i,:)),[1/13.5 1/11.5], 1); 
  scaled_eq_z.V_highpass(i,:)  = highpass(squeeze(scaled_eq_z.V(i,:)),1/11.5, 1);   
  scaled_eq_z.V_lowpass(i,:)   = lowpass(squeeze(scaled_eq_z.V(i,:)),1/11.5, 1);   
  scaled_eq_z.V_inertial(i,:)  = bandpass(squeeze(scaled_eq_z.V(i,:)),[1/inertial_end 1/inertial_start], 1); 
end

%% Kinetic energy PSD 
NFFT = length(scaled_eq_z.U(:,1));
win = hann(NFFT);
Fs = 1/(abs(z_stretched_equal_spaced(2))-abs(z_stretched_equal_spaced(1))); 
waveno = 0:Fs/NFFT:Fs/2; % we can do this because the depth

clear psd psdx_array
[ke.psd_total, psdx_array] = welch_estimate_corrected_rt(scaled_eq_z.U,scaled_eq_z.V,win,NFFT,Fs); 
[ke.psd_6_8hr, ~]          = welch_estimate_corrected_rt(scaled_eq_z.U_6_to_8hr,scaled_eq_z.V_6_to_8hr,win,NFFT,Fs); 
[ke.psd_sem_diur, ~]       = welch_estimate_corrected_rt(scaled_eq_z.U_sem_diur,scaled_eq_z.V_sem_diur,win,NFFT,Fs); 
[ke.psd_highpass, ~]       = welch_estimate_corrected_rt(scaled_eq_z.U_highpass,scaled_eq_z.V_highpass,win,NFFT,Fs); 
[ke.psd_lowpass, ~]        = welch_estimate_corrected_rt(scaled_eq_z.U_lowpass,scaled_eq_z.V_lowpass,win,NFFT,Fs); 
[ke.psd_inertial, ~]       = welch_estimate_corrected_rt(scaled_eq_z.U_inertial,scaled_eq_z.V_inertial,win,NFFT,Fs); 

ke.waveno       = waveno;
ke.psd_total    = ke.psd_total./(NFFT*Fs);
ke.psd_6_8hr    = ke.psd_6_8hr./(NFFT*Fs);
ke.psd_sem_diur = ke.psd_sem_diur./(NFFT*Fs);
ke.psd_lowpass  = ke.psd_lowpass./(NFFT*Fs);
ke.psd_highpass = ke.psd_highpass./(NFFT*Fs);
ke.psd_inertial = ke.psd_inertial./(NFFT*Fs);
psd_array_normalised = psdx_array./(NFFT*Fs);
 

%% bootstrap confidence intervals. The following is taken from b028 routine of my GRL paper
% alpha=0.1 below is for 90% confidence interval. For 95% use alpha=0.05
% (this is also deafault so you may not use it too)

%bootci(nboot,@(x)[mean(x) std(x)],y); nboot need not be the sample length.
%it has has to be the number of samples we need to form. the length of each
%sample is always the length of y formed with replacement.

% I have tested both the commands below and they give same answer (to second
% decimal place). (01 April 2022)
%for j = 1:length(psdx_array(1,:))
% psd_ci(:,j) = bootci(length(psd_array_normalised(:,1)),{@mean,psd_array_normalised(:,j)},'alpha',0.05,'type','norm');
%end

for j = 1:length(psdx_array(1,:))
 psd_ci(:,j) = bootci(100,{@(x)mean(x),psd_array_normalised(:,j)},'alpha',0.05,'type','norm');
end

ke.psd_ci = psd_ci;


%% Pinkel estimate of shear

NFFT_waveno_shear=length((scaled_eq_z.U(:,1))); %taking FFT along the whole vertical grid.
% I am subtracting 1 because while taking derivative in calculating shear,
% we have one less point in the vertical (22 Jan 2022)
%delta_z = 2*pi/(abs(z_stretched_equal_spaced(2))-abs(z_stretched_equal_spaced(1))); 
%delta_z_unstretched = 2*pi/(abs(z_unstretched_equal(2)-z_unstretched_equal(1)));
% removed the 2pi in the wavenumber following early literature (Gregg et
% al 1991: "varieties of fully resolved spectra of vertical shear"): 4Nov21
delta_z = 1/(abs(z_stretched_equal_spaced(2)-z_stretched_equal_spaced(1))); 
% to have it in cycles per km, multiply by thousand
% the following was used to see if we get different spectra when using an
% unstretched coordinate. It obviously does not change the spectra
%delta_z_unstretched = 1/(abs(z_unstretched_equal(2)-z_unstretched_equal(1)));

% These are the wavenumbers.  we can do this because the depth
% array is equally spaced albeit in stretched coordinates. This is
% different when we calculate both freq-waveno because there we create
% wavenos from -ve to +ve delta_z/2. The one is subtracted because shear
% has different wavenumbers than velocities
waveno_shear = 0:delta_z/(NFFT_waveno_shear-1):delta_z/2;
%waveno_shear_unstretched = 0:delta_z_unstretched/NFFT_waveno_shear:delta_z_unstretched/2;

clear psdx_array_shear wave_freq_psd_avg_shear psdx_array_shear_unstretched wave_psd_avg_shear_unstretched
%%%%%%%%%%%%%%%%
% calling the waveno calculation routine. Complex velocity is calculated
% inside the following routine. This is a single-sided spectra and hence
% has a factor 2 multiplied inside routine
[psdx_array_shear,wave_psd_avg_shear] = waveno_only_spectra_shear_ybp(scaled_eq_z.U,scaled_eq_z.V,z_stretched_equal_spaced,NFFT_waveno_shear);
%%%%%%%%%%%%%%%%

% psd in different bands
[~,wave_psd_avg_shear_6_8hr]    = waveno_only_spectra_shear_ybp(scaled_eq_z.U_6_to_8hr,scaled_eq_z.V_6_to_8hr,z_stretched_equal_spaced,NFFT_waveno_shear);
[~,wave_psd_avg_shear_sem_diur] = waveno_only_spectra_shear_ybp(scaled_eq_z.U_sem_diur,scaled_eq_z.V_sem_diur,z_stretched_equal_spaced,NFFT_waveno_shear);
[~,wave_psd_avg_shear_highpass] = waveno_only_spectra_shear_ybp(scaled_eq_z.U_highpass,scaled_eq_z.V_highpass,z_stretched_equal_spaced,NFFT_waveno_shear);
[~,wave_psd_avg_shear_inertial] = waveno_only_spectra_shear_ybp(scaled_eq_z.U_inertial,scaled_eq_z.V_inertial,z_stretched_equal_spaced,NFFT_waveno_shear);
% Normalising the psd with the wavenumber. Shear and strain both are
% normalised by win^2/Length_of_window to account for the lost variance due
% to windowing
wave_psd_avg_shear = wave_psd_avg_shear./((NFFT_waveno_shear-1)*delta_z);
wave_psd_avg_shear_6_8hr = wave_psd_avg_shear_6_8hr./((NFFT_waveno_shear-1)*delta_z);
wave_psd_avg_shear_sem_diur = wave_psd_avg_shear_sem_diur./((NFFT_waveno_shear-1)*delta_z);
wave_psd_avg_shear_highpass = wave_psd_avg_shear_highpass./((NFFT_waveno_shear-1)*delta_z);
wave_psd_avg_shear_inertial = wave_psd_avg_shear_inertial./((NFFT_waveno_shear-1)*delta_z);
%wave_psd_avg_shear_unstretched = wave_psd_avg_shear_unstretched/(NFFT_waveno_shear*delta_z_unstretched); 
 
psd_shear.pinkelwaveno         = waveno_shear;  
psd_shear.greggwaveno          = waveno;  
psd_shear.pinkel_shear_psd     = wave_psd_avg_shear;
psd_shear.pinkel_shear_6_8hr   = wave_psd_avg_shear_6_8hr;
psd_shear.pinkel_shear_sem_diur= wave_psd_avg_shear_sem_diur;
psd_shear.pinkel_shear_highpass= wave_psd_avg_shear_highpass;
psd_shear.pinkel_shear_inertial= wave_psd_avg_shear_inertial;
psd_shear.comment = 'the spectra are averaged for 73 days from 1 March 2006 in the depth range 80-1400m';
 
%% Gregg estimate of shear
psd_shear.gregg_shear_psd      = (2*pi.*ke.waveno).^2.*ke.psd_total; 
psd_shear.gregg_shear_6_8hr    = (2*pi.*ke.waveno).^2.*ke.psd_6_8hr; 
psd_shear.gregg_shear_sem_diur = (2*pi.*ke.waveno).^2.*ke.psd_sem_diur; 
psd_shear.gregg_shear_lowpass  = (2*pi.*ke.waveno).^2.*ke.psd_lowpass;
psd_shear.gregg_shear_highpass = (2*pi.*ke.waveno).^2.*ke.psd_highpass;
psd_shear.gregg_shear_inertial = (2*pi.*ke.waveno).^2.*ke.psd_inertial; 

end
end

